//
// Created by tong on 9/28/24.
//

#include "MeshInterpolation.h"
#include <iostream>


bool MeshInterpolation::perfFlag_ = true;
double MeshInterpolation::maxRadius_ = 100;
unsigned int MeshInterpolation::threadNum_ = 1;
bool MeshInterpolation::perfTestMode_ = false;

void MeshInterpolation::buildInterpolationStencil(const KDTreeL2Type &srcTree, const PointCloud &destPoints, InterpolationStencil &destStencil)
{   
    const size_t destPointsSize = destPoints.size_;
    const size_t srcPointsSize = srcTree.size_;

    const double sizeRatio = double(destPointsSize) / double(srcPointsSize);
    
    int stencilSize =1;
    double IDWExp = 1;

    if (sizeRatio>200)
    {
        stencilSize = 8;
        IDWExp = 0.5;
    }
    else if (sizeRatio <= 200 && sizeRatio > 10)
    {
        stencilSize = 6;
        IDWExp = 0.5;
    }
    else if(sizeRatio <= 10 && sizeRatio > 5)
    {
        stencilSize = 6;
        IDWExp = 1;
    }
    else if(sizeRatio <= 5 && sizeRatio > 2)
    {
        stencilSize = 5;
        IDWExp = 2;
    }
    else if(sizeRatio <= 2 && sizeRatio > 1) {
        stencilSize = 3;
        IDWExp = 4;
    }
    else if(sizeRatio <= 1 && sizeRatio > 0.5) {
        stencilSize = 2;
        IDWExp = 6;
    }
    else if(sizeRatio<=0.5 && sizeRatio >0.1 )
    {
        stencilSize = 2; 
        IDWExp = 6;
    }
    else if(sizeRatio<=0.1 && sizeRatio >0.005)
    {
        stencilSize = 2;
        IDWExp = 8;
    }
    else if(sizeRatio<=0.005)
    {
        stencilSize = 1;
        IDWExp = 1;
    }

    std::cout << "Size ratio: " << sizeRatio << ", stencil size: " << stencilSize << ", IDWExp: " << IDWExp << std::endl;
    LabelListList neighbors(destPointsSize);
    DoubleListList distance(destPointsSize);
    
    // search for the nearest neighbors
    //std::cout<< "destPoints size: "<<srcTree.leaf_max_size_ << std::endl;
    for (size_t pointI = 0; pointI < destPointsSize; pointI++)
    {
        const Point &dPoint = destPoints[pointI];
        LabelList resultIdx(stencilSize, 0);
        DoubleList resultDist(stencilSize, 0.0);
        KNNResultType resultSet(stencilSize);
        resultSet.init(&resultIdx[0], &resultDist[0]);
        double pointCoord[3] = {dPoint.x, dPoint.y, dPoint.z};
        srcTree.findNeighbors(resultSet, &pointCoord[0]);
        if (!resultSet.full())
        {
            std::cout << "Point: " << pointI << " not find enough neighbors!" << std::endl;
            std::abort();
        }
        neighbors[pointI] = resultIdx;
        distance[pointI] = resultDist;
    }

    // calculate the weight
    destStencil.indicator.push_back(0);
    for (size_t pointI = 0; pointI < destPointsSize; pointI++)
    {
        const LabelList &neiIdx = neighbors[pointI];
        const DoubleList &neiDist = distance[pointI];
        double dSum = 0;
        size_t stencilCounter = 0;
        for (int neiI = 0; neiI < stencilSize; neiI++)
        {
            if(neiDist[neiI] < 1e-6) {
                destStencil.stencil.push_back(int(neiIdx[neiI]));
                destStencil.weight.push_back(1.0);
                dSum = 1.0;
                ++stencilCounter;
                break;
            }
            destStencil.stencil.push_back(int(neiIdx[neiI]));
            destStencil.weight.push_back(1.0 / pow(neiDist[neiI],IDWExp));
            dSum += destStencil.weight.back();
            ++stencilCounter;
        }
        destStencil.indicator.push_back(int(stencilCounter)+destStencil.indicator.back());
        for (int neiI = destStencil.indicator[pointI]; neiI < destStencil.indicator[pointI+1]; neiI++)
        {
            destStencil.weight[neiI] /= dSum;
        }
    }
}

MeshInterpolation::MeshInterpolation(const PointList &mesh1Points, const PointList &mesh2Points)
    : mesh1PointCloud_(mesh1Points,mesh1Points.size()),
      mesh2PointCloud_(mesh2Points,mesh2Points.size()),
      mesh1KDTreePtr_(nullptr),
      mesh2KDTreePtr_(nullptr)
{
    clock_t startT = std::clock();
    mesh1KDTreePtr_.reset(new KDTreeL2Type(3, mesh1PointCloud_));
    clock_t endT1 = std::clock();
    if (perfFlag_)
    {
        std::cout << "Time spent in build mesh1 KDTree: " << double(endT1 - startT) / CLOCKS_PER_SEC << " s" << std::endl;
    }
    mesh2KDTreePtr_.reset(new KDTreeL2Type(3, mesh2PointCloud_));
    clock_t endT2 = std::clock();
    if (perfFlag_)
    {
        std::cout << "Time spent in build mesh2 KDTree: " << double(endT2 - endT1) / CLOCKS_PER_SEC << " s" << std::endl;
    }

    // build mesh 1 to mesh 2 stencil
    buildInterpolationStencil(*mesh1KDTreePtr_, mesh2PointCloud_, mesh1ToMesh2IpData_);
    clock_t endT3 = std::clock();
    if (perfFlag_)
    {
        std::cout << "Time spent in build mesh1 stencil: " << double(endT3 - endT2) / CLOCKS_PER_SEC << " s" << std::endl;
    }

    buildInterpolationStencil(*mesh2KDTreePtr_, mesh1PointCloud_, mesh2ToMesh1IpData_);
    clock_t endT4 = std::clock();
    if (perfFlag_)
    {
        std::cout << "Time spent in build mesh2 stencil: " << double(endT4 - endT3) / CLOCKS_PER_SEC << " s" << std::endl;
    }
}

MeshInterpolation::MeshInterpolation(const PointList &mesh1Points, const PointList &mesh2Points, const size_t mesh1Size, const size_t mesh2Size)
    : mesh1PointCloud_(mesh1Points,mesh1Size),
      mesh2PointCloud_(mesh2Points,mesh2Size),
      mesh1KDTreePtr_(nullptr),
      mesh2KDTreePtr_(nullptr)
{
    clock_t startT = std::clock();
    mesh1KDTreePtr_.reset(new KDTreeL2Type(3, mesh1PointCloud_));
    clock_t endT1 = std::clock();
    if (perfFlag_)
    {
        std::cout << "Time spent in build mesh1 KDTree: " << double(endT1 - startT) / CLOCKS_PER_SEC << " s" << std::endl;
    }
    mesh2KDTreePtr_.reset(new KDTreeL2Type(3, mesh2PointCloud_));
    clock_t endT2 = std::clock();
    if (perfFlag_)
    {
        std::cout << "Time spent in build mesh2 KDTree: " << double(endT2 - endT1) / CLOCKS_PER_SEC << " s" << std::endl;
    }

    // build mesh 1 to mesh 2 stencil
    buildInterpolationStencil(*mesh1KDTreePtr_, mesh2PointCloud_, mesh1ToMesh2IpData_);
    clock_t endT3 = std::clock();
    if (perfFlag_)
    {
        std::cout << "Time spent in build mesh1 stencil: " << double(endT3 - endT2) / CLOCKS_PER_SEC << " s" << std::endl;
    }

    buildInterpolationStencil(*mesh2KDTreePtr_, mesh1PointCloud_, mesh2ToMesh1IpData_);
    clock_t endT4 = std::clock();
    if (perfFlag_)
    {
        std::cout << "Time spent in build mesh2 stencil: " << double(endT4 - endT3) / CLOCKS_PER_SEC << " s" << std::endl;
    }
}

void MeshInterpolation::checkDataSize(const DoubleList &src, const InterpolationStencil &stencil, DoubleList &dest,const size_t dest_size)
{

    if (dest_size != (stencil.indicator.size() - 1) || dest.size() < dest_size)
    {
        std::cout << "Inconsistent stencil! destination field size=" << dest.size()
                  << ", stencil size= " << stencil.indicator.size() - 1 << std::endl;
        std::abort();
    }

    int maxIdx = 0;
    for (size_t pointI = 0; pointI < stencil.stencil.size(); pointI++)
    {
        maxIdx = max(maxIdx, stencil.stencil[pointI]);
    }
    if (maxIdx > int(src.size()))
    {
        std::cout << "Stencil data outbound! source data size="
                  << src.size() << ", Max stencil index=" << maxIdx << std::endl;
        std::abort();
    }
}

void MeshInterpolation::interpolate (const DoubleList& src,const InterpolationStencil& stencil,DoubleList& dest,const size_t destSize)
{

    checkDataSize(src, stencil, dest, destSize);
    const IntList &indicator = stencil.indicator;
    const IntList &pStencil = stencil.stencil;
    const DoubleList &weight = stencil.weight;
    const int stencilSize = int(destSize);
    auto ipFunc = [&src, &dest, &indicator, &pStencil, &weight](const int startPos, const int endPos)
    {   
        // for perf test loop 10000 times
        if (perfTestMode_)
        {
            for (int r = 0; r < 10000; r++)
            {
                for (int i = startPos; i < endPos; ++i)
                {
                    dest[i] = 0;
                    const int startIdx = indicator[i];
                    const int endIdx = indicator[i + 1];
                    for (int idx = startIdx; idx < endIdx; ++idx)
                    {
                        dest[i] += src[pStencil[idx]] * weight[idx];
                    }
                }
            }
        }
        else // normal mode only loop once
        {
            for (int i = startPos; i < endPos; ++i)
            {
                dest[i] = 0;
                const int startIdx = indicator[i];
                const int endIdx = indicator[i + 1];
                for (int idx = startIdx; idx < endIdx; ++idx)
                {
                    dest[i] += src[pStencil[idx]] * weight[idx];
                }
            }
        }
    };
    if (threadNum_ > 1)
    {
        vector<std::thread> threads;
        const int segment = stencilSize / threadNum_;
        for (unsigned int i = 0; i < threadNum_; ++i)
        {
            threads.emplace_back(ipFunc, i * segment, i == (threadNum_ - 1) ? stencilSize : (i + 1) * segment);
        }
        auto startT = std::chrono::high_resolution_clock::now();
        for (auto &t : threads)
        {
            t.join();
        }
        auto endT = std::chrono::high_resolution_clock::now();
        auto duration1 = std::chrono::duration_cast<std::chrono::microseconds>(endT - startT);
        if (perfFlag_)
        {
            std::cout << "Time spent in interpolation: " << duration1.count() / 1e6 << " s" << std::endl;
        }
    }
    else
    {
        auto startT = std::chrono::high_resolution_clock::now();
        ipFunc(0, stencilSize);
        auto endT = std::chrono::high_resolution_clock::now();
        auto duration1 = std::chrono::duration_cast<std::chrono::microseconds>(endT - startT);
        if (perfFlag_)
        {
            std::cout << "Time spent in interpolation: " << duration1.count() / 1e6 << " s" << std::endl;
        }
    }
    // for (int i = 0; i < stencilSize; ++i)
    // {
    //     const int startIdx = indicator[i];
    //     const int endIdx = indicator[i + 1];
    //     for (int idx = startIdx; idx < endIdx; ++idx)
    //     {
    //         dest[i] += src[pStencil[idx]] * weight[idx];
    //     }
    // }
}

void MeshInterpolation::interpolateMesh1toMesh2(const DoubleList &src, DoubleList &dest,const size_t destSize) const
{
    interpolate(src,mesh1ToMesh2IpData_,dest,destSize);
}
void MeshInterpolation::interpolateMesh1toMesh2(const DoubleList &src, DoubleList &dest) const
{
    interpolate(src,mesh1ToMesh2IpData_,dest,dest.size());
}
void MeshInterpolation::interpolateMesh2toMesh1(const DoubleList &src, DoubleList &dest,const size_t destSize) const
{
    interpolate(src,mesh2ToMesh1IpData_,dest,destSize);
}
void MeshInterpolation::interpolateMesh2toMesh1(const DoubleList &src, DoubleList &dest) const
{
    interpolate(src,mesh2ToMesh1IpData_,dest,dest.size());
}


DoubleList calcError(const DoubleList& srcField,const DoubleList& desField) {
    if(srcField.size()!=desField.size()) {
        std::cerr<<"Two Field size not match!"<<std::endl;
        std::abort();
    }
    DoubleList norDiff(srcField.size());
    for(size_t i = 0; i<srcField.size();i++) {
        norDiff[i]=abs(srcField[i]-desField[i])/abs(srcField[i]);
    }
    double maxError = 0;
    double avgError = 0;
    for(size_t i = 0; i<norDiff.size();i++) {
        maxError = max(maxError,norDiff[i]);
        avgError += norDiff[i];
    }
    avgError /= norDiff.size();

    std::cout << "Max Error: " << maxError * 100 << "%, Average Error: " << avgError * 100 << "%" << std::endl;

    return norDiff;
}


void MeshInterpolation::freeTree()
{
    mesh1KDTreePtr_.reset();
    mesh2KDTreePtr_.reset();
}
void MeshInterpolation::freeAll()
{
    mesh1KDTreePtr_.reset();
    mesh2KDTreePtr_.reset();
    mesh1ToMesh2IpData_.stencil.clear();
    mesh1ToMesh2IpData_.weight.clear();
    mesh1ToMesh2IpData_.indicator.clear();
    mesh2ToMesh1IpData_.stencil.clear();
    mesh2ToMesh1IpData_.weight.clear();
    mesh2ToMesh1IpData_.indicator.clear();
}